import torch
import cv2
import numpy as np
from NumberNet import NumberNet

model = torch.load("numbers_model1.pth", weights_only=False)

img = cv2.imread("./data/test/8.png")
img = np.expand_dims(img, 0)  # 1[0]x28[1]x28[2]x3[3]
img = torch.from_numpy(img)
# 直接预测结果(torch的tensor在datasets加载下不可以使用reshape进行图像更改)
img = torch.permute(img, [0, 3, 1, 2])  # 1x3x28x28
print(img.shape)
model = model.to(torch.device("cpu"))
predict = model(img.float())
result = torch.argmax(predict, dim=-1)
print(result[0].item())
